import copy
import os
import pickle
import shutil
import random

import click
import torch
from torch import optim
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from segmentation.dataset import US
from segmentation.utils import *

def train(
    n_train=3,

    split=[72,8,20],

    data_dir='data/US',

):

    split_length = len(os.listdir(os.path.join(data_dir, 'images')))
    split = [int(s) for s in split]
    split = [split_length*s//sum(split) for s in split]
    split[1] = split_length - split[0] - split[2]
    print(f'dataset split: {split}')
    for d in [
        'save_seg/supervised/exp16-222shot-reso256-augTrue/checkpoint'

    ]:
        sum_evals = 0
        print(d.split('/')[-2])
        for i_train in range(n_train):
            random.seed(28*i_train)

            # train_loader, val_loader, test_loader = get_dataloader(shot, batch_size, dataset, split, i_train)
            samples = list(range(split_length))
            random.shuffle(samples)
            test_loader = DataLoader(
                US(
                    data_dir=data_dir,
                    aug=False,
                    sample=samples[split[0]+split[1]:split_length],
                ),
                batch_size=5, shuffle=False, num_workers=4, pin_memory=True, drop_last=False
            )

            # test part of i train
            with torch.no_grad():
                with open(os.path.join(d, f'{i_train}_best.pth'), 'rb') as f:
                    net = pickle.load(f)['net'].eval().requires_grad_(False).cuda()

                evaluations = eval_dice_2d(net, test_loader, False, '', i_train)
                sum_evals += evaluations
            
            for i, e in enumerate(['dice', 'accu', 'iou', 'spec', 'sens', 'ppv', 'npv']):
                print(f'{evaluations[i].item()*100:.1f}, ', end='')
            print()

        evaluations = sum_evals/n_train
        for i, e in enumerate(['dice', 'accu', 'iou', 'spec', 'sens', 'ppv', 'npv']):
            print(f'{evaluations[i].item()*100:.1f}, ', end='')
        print()

if __name__ == "__main__":
    train()
#----------------------------------------------------------------------------
